import os
from dotenv import load_dotenv, find_dotenv
from langchain_community.chat_models import ChatZhipuAI 
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain.callbacks.streaming_stdout import (
    StreamingStdOutCallbackHandler
)

_ = load_dotenv(find_dotenv())
model = os.environ.get('model')
if model is None:
    raise ValueError("model is not set in the .env file")

prompt_template = ChatPromptTemplate.from_template(
    "请写一首关于{topic}的诗，要求不超过{num}行。"
    )
chat_model = ChatZhipuAI(model=model,
                  temperature=0.9,    
                  callbacks=[StreamingStdOutCallbackHandler()],       
    )
output_parser = StrOutputParser()

chain = prompt_template | chat_model | output_parser

response = chain.stream({"topic":"春天", "num":8})
for chunk in response:
   pass